import random
import threading
import time
import zmq
from kvsimple import KVMsg
from zhelpers import zpipe


def main():
    ctx = zmq.Context()
    publisher = ctx.socket(zmq.PUB)
    publisher.bind("tcp://*:5557")

    updates, peer = zpipe(ctx)

    manager_thread = threading.Thread(target=state_manager, args=(ctx, peer))
    manager_thread.daemon = True
    manager_thread.start()

    sequence = 0
    random.seed(time.time())

    try:
        while True:
            sequence += 1
            kvmsg = KVMsg(sequence)
            kvmsg.key = "%d" % random.randint(1, 10000)
            kvmsg.body = "%d" % random.randint(1, 1000000)
            kvmsg.send(publisher)
            kvmsg.send(updates)
    except KeyboardInterrupt:
        print("Interrupted %d messages out" % sequence)


def state_manager(ctx, pipe):
    kvmap = {}
    pipe.send(b"READY")
    snapshot = ctx.socket(zmq.ROUTER)
    snapshot.bind("tcp://*:5556")

    poller = zmq.Poller()
    poller.register(pipe, zmq.POLLIN)
    poller.register(snapshot, zmq.POLLIN)

    sequence = 0
    while True:
        try:
            items = dict(poller.poll())
        except (zmq.ZMQError, KeyboardInterrupt):
            break

        if pipe in items:
            kvmsg = KVMsg.recv(pipe)
            sequence = kvmsg.sequence
            kvmsg.store(kvmap)
        if snapshot in items:
            msg = snapshot.recv_multipart()
            identity = msg[0]
            request = msg[1]
            if request == b"ICANHAZ?":
                pass
            else:
                print("E: bad request, aborting")
                break
            route = Route(snapshot, identity)
            for k, v in kvmap.items():
                send_single(k, v, route)

            print("Sending state shapshot=%s" % sequence)
            snapshot.send(identity, zmq.SNDMORE)
            kvmsg = KVMsg(sequence)
            kvmsg.key = "KTHXBAI"
            kvmsg.body = ""
            kvmsg.send(snapshot)


class Route(object):
    def __init__(self, socket, identity):
        self.socket = socket
        self.identity = identity


def send_single(key, kvmsg, route):
    route.socket.send(route.identity, zmq.SNDMORE)
    kvmsg.send(route.socket)


if __name__ == '__main__':
    main()
